package alloy.transform.atomize;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;

import alloy.ast.ArrowMultExpr;
import alloy.ast.Expr;
import alloy.ast.Multiplicity;
import alloy.ast.MultiplicityExpr;
import alloy.ast.SetMultExpr;
import alloy.ast.SigmaExpr;
import alloy.util.Dbg;

/**
 * A helper class for atomizing MultExpr. The fundamental issue it tries to
 * tackle is the fact that for an ArrowMultExpr there can be infinitely many
 * levels of ArrowMultExpr in the left/right expressions.
 */
class MultExprTermsIterator implements Iterator {

	private class MultExprBoolean {
		boolean bool;

		MultiplicityExpr multExpr;
	}

	/**
	 * the original MultExpr for which we are generating atomized subexprs
	 */
	private MultiplicityExpr oldMultExpr;

	/**
	 * maps a MultExpr in the tree of the original MultExpr to a List of its
	 * atomized subexprs. Note that only the leaves of the tree (i.e. only
	 * SetMultExpr's) will be mapped.
	 */
	private Map multExprToTerms;

	/**
	 * maps a MultExpr in the tree of the original MultExpr to the index of the
	 * current atomized subexpr
	 */
	private Map multExprToIndex;

	private boolean first;

	private boolean finished;

	public MultExprTermsIterator(final MultiplicityExpr multExpr,
			final AtomizationVisitor av) {
		oldMultExpr = multExpr;
		multExprToTerms = new LinkedHashMap();
		multExprToIndex = new LinkedHashMap();
		finished = false;
		first = true;

		// perform the first walk down the tree
		initialize(multExpr, av);
	}

	private void initialize(final MultiplicityExpr multExpr,
			final AtomizationVisitor av) {

		if (multExpr instanceof ArrowMultExpr) {
			// split into the left and right exprs
			MultiplicityExpr leftMultExpr = ((ArrowMultExpr) multExpr)
					.getLeft();
			initialize(leftMultExpr, av);

			MultiplicityExpr rightMultExpr = ((ArrowMultExpr) multExpr)
					.getRight();
			initialize(rightMultExpr, av);

		} else {
			// A SetMultExpr -- no need to split, just atomize straight up
			SigmaExpr sigmaExpr = (SigmaExpr) multExpr.toExpr()
					.applyReturnVisitor(av);
			multExprToTerms.put(multExpr, new ArrayList(sigmaExpr.getTerms()));
			multExprToIndex.put(multExpr, new Integer(0));
		}

	}

	public void remove() {
		throw new UnsupportedOperationException();
	}

	public Object next() {
		// traverse tree

		if (finished) {
			// no more elements exist
			throw new NoSuchElementException();

		} else if (first) {
			// first element to return, do not need to increment
			MultExprBoolean ret = traverse(oldMultExpr, false);
			first = false;
			if (indicesAtMax()) {
				// this is the last element
				finished = true;
			}
			return ret.multExpr;

		} else {
			// normal case, increment
			MultExprBoolean ret = traverse(oldMultExpr, true);
			if (ret.bool) {
				if (indicesAtMax()) {
					// this is the last element
					finished = true;
				}
				return ret.multExpr;
			} else {
				Dbg
						.fail("Unexpected error in MultExprTermsIter -- unsuccesful increment in next()");
				return null;
			}
		}
	}

	/**
	 * traverses the tree of MultExpr, incrementing the first index seen by one.
	 * The order of traversal is deterministic, so this works.
	 * 
	 * @param increment
	 *            whether or not we need to increment an index
	 * @return boolean true iff an increment was successful
	 */
	private MultExprBoolean traverse(MultiplicityExpr multExpr,
			boolean increment) {

		MultExprBoolean ret = new MultExprBoolean();

		if (multExpr instanceof ArrowMultExpr) {

			ArrowMultExpr arrowMult = (ArrowMultExpr) multExpr;
			// split into the left and right exprs

			// recurse on left
			MultiplicityExpr leftMultExpr = arrowMult.getLeft();

			MultExprBoolean leftRet = traverse(leftMultExpr, increment);

			// recurse on right--only need to increment if the original
			// increment is true
			// and incrementing was unsuccessful in left
			MultiplicityExpr rightMultExpr = arrowMult.getRight();
			MultExprBoolean rightRet = traverse(rightMultExpr,
					(increment && !leftRet.bool));

			// the return boolean is true if either succeeded
			ret.bool = leftRet.bool || rightRet.bool;

			// construct the new MultExpr
			// not copying leftRet.multExpr, etc., because it is assumed to have
			// been newly constructed
			ArrowMultExpr newMultExpr = new ArrowMultExpr(arrowMult
					.getLocation(), (Multiplicity) arrowMult.getMult().copy(),
					leftRet.multExpr, rightRet.multExpr);
			newMultExpr.setType(leftRet.multExpr.getType().product(
					rightRet.multExpr.getType()));
			ret.multExpr = newMultExpr;

		} else {
			// A SetMultExpr -- base case

			Dbg.chk(multExprToTerms.containsKey(multExpr));
			Dbg.chk(multExprToIndex.containsKey(multExpr));

			List terms = (List) multExprToTerms.get(multExpr);
			int index = ((Integer) multExprToIndex.get(multExpr)).intValue();

			// if index is not at the max and we need to increment, do so now
			if (increment) {

				if (index < terms.size() - 1) {

					// able to just increment this index
					index++;
					multExprToIndex.put(multExpr, new Integer(index));

					// successful
					ret.bool = true;

				} else {

					// this index is at its max, so we're gonna have to
					// increment the next one.
					// we assume that eventually we'll encoutner one that's
					// incrementable because
					// otherwise, traverse should not have been called. With
					// this assumption,
					// we can safely reset this index to 0, so that we get
					// patterns of the following
					// form:
					// 0 0 0
					// 1 0 0
					// 2 0 0
					// 0 1 0
					// 1 1 0
					// 2 1 0
					// 0 2 0
					// 1 2 0
					// 2 2 0
					// ...

					index = 0;
					multExprToIndex.put(multExpr, new Integer(index));
					ret.bool = false;

				}

			} else {
				ret.bool = false;
			}

			Expr currentTerm = (Expr) terms.get(index);

			MultiplicityExpr newMultExpr = new SetMultExpr(multExpr
					.getLocation(), (Multiplicity) ((SetMultExpr) multExpr)
					.getMult().copy(), (Expr) currentTerm.copy());
			newMultExpr.setType(currentTerm.getType());

			ret.multExpr = newMultExpr;
		}

		return ret;

	}

	private boolean indicesAtMax() {
		// check if all indices are set to their maximum

		for (Iterator exprIter = multExprToTerms.keySet().iterator(); exprIter
				.hasNext();) {
			MultiplicityExpr currentExpr = (MultiplicityExpr) exprIter.next();
			List terms = (List) multExprToTerms.get(currentExpr);
			int index = ((Integer) multExprToIndex.get(currentExpr)).intValue();
			if (index < terms.size() - 1) {
				return false;
			}
		}

		return true;
	}

	public boolean hasNext() {
		return !finished;
	}
}
